/* Copyright (C) 2011 clueless <clueless.mmopeon@gmail.com>
 *
 * This file is part of MMO Peon.
 *
 * MMO Peon is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MMO Peon is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MMO Peon.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <wx/wx.h>
#include <windows.h>
#include <stddef.h>

// TODO: - lots of error checking
//       - Freeing up resources


/*
CPU Disasm
Address   Hex dump          Command
00000000  90            NOP
00000001  90            NOP
00000002  68 EFBEADDE   PUSH DEADBEEF       Push address of DLL path
00000007  BA EFBEADDE   MOV EDX,DEADBEEF    move address of GetModuleHandle in EDX
0000000C  FFD2          CALL EDX            call it
0000000E  50            PUSH EAX            push return value of GetModuleHandle
0000000F  BA EFBEADDE   MOV EDX,DEADBEEF    grab the address of FreeLibrary
00000014  FFD2          CALL EDX            call it...
00000016  C20400        RET 0x04
00000019  00                                Filler
*/

// EB FE
const char *ejectStub = "\x90\x90"
                        "\x68\xEF\xBE\xAD\xDE"
                        "\xBA\xEF\xBE\xAD\xDE"
                        "\xFF\xD2"
                        "\x50"
                        "\xBA\xEF\xBE\xAD\xDE"
                        "\xFF\xD2"
                        "\xC2\x04\x00"
                        "\x00";

const DWORD PROCESS_INJECT_DLL = PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION |
                                 PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ;

struct EjectStub
{
    char        res1_[3];
    DWORD_PTR   dllString;
    char        res2_;
    DWORD_PTR   getModuleHandle;
    char        res3_[4];
    DWORD_PTR   freeLibrary;
    char        res5_[6];
    char        extraData[1];
} __attribute__((packed));


HANDLE InjectorOpenProcess(const char *pid)
{
    HANDLE proc = OpenProcess(PROCESS_INJECT_DLL, 0, atol(pid));
    if (proc == NULL)
        wxMessageBox(wxT("Failed to open process"), wxT("Error!"), wxOK | wxICON_ERROR);

    return proc;
}

void *InjectorAlloc(HANDLE proc, DWORD size, DWORD protection)
{
    void *page = VirtualAllocEx(proc, NULL, size, MEM_COMMIT, protection);
    if (page == NULL)
        wxMessageBox(wxT("Failed to allocate remote memory"), wxT("Error!"), wxOK | wxICON_ERROR);

    return page;
}

bool InjectorWPM(HANDLE proc, void *dest, const void *source, DWORD size)
{
    if (!WriteProcessMemory(proc, dest, source, size, NULL))
    {
        wxMessageBox(wxT("Failed to write process memory"), wxT("Error!"), wxOK | wxICON_ERROR);
        return false;
    }
    return true;
}

bool InjectorCreateThread(HANDLE proc, LPTHREAD_START_ROUTINE procedure, void *argument)
{
    if (!CreateRemoteThread(proc, NULL, 0, procedure, argument, 0, NULL))
    {
        wxMessageBox(wxT("Failed create remote thread"), wxT("Error!"), wxOK | wxICON_ERROR);
        return false;
    }
    return true;
}

// Eject a dll from target process
int DllEject(const char *pid, const char *dllName)
{
    HANDLE proc = InjectorOpenProcess(pid);
    int size = sizeof(struct EjectStub) + strlen(dllName) + 1;
    void *page = InjectorAlloc(proc, size, PAGE_EXECUTE_READWRITE);
    DWORD_PTR gmhAddress = (DWORD_PTR)GetProcAddress(GetModuleHandle(L"kernel32.dll"), "GetModuleHandleA");
    DWORD_PTR flAddress = (DWORD_PTR)GetProcAddress(GetModuleHandle(L"kernel32.dll"), "FreeLibrary");

    EjectStub *cs = (EjectStub *)malloc(size);
    memcpy(cs, ejectStub, sizeof(EjectStub));
    strcpy(cs->extraData, dllName);

    cs->getModuleHandle = (uintptr_t)gmhAddress;
    cs->freeLibrary = (uintptr_t)flAddress;
    cs->dllString = (uintptr_t)page + offsetof(struct EjectStub, extraData);
    InjectorWPM(proc, page, cs, size);
    DWORD lol;
    VirtualProtectEx(proc, page, size, PAGE_EXECUTE_READ, &lol);
    InjectorCreateThread(proc, (LPTHREAD_START_ROUTINE)page, NULL);
    return 0;
}

// Inject a DLL into the target process
int DllInject(const char *pid, const char *dllPath)
{
    HANDLE proc = InjectorOpenProcess(pid);
    void *page = InjectorAlloc(proc, strlen(dllPath) + 1, PAGE_READWRITE);
    InjectorWPM(proc, page, dllPath, strlen(dllPath) + 1);
    LPTHREAD_START_ROUTINE loadLibraryAddress = (LPTHREAD_START_ROUTINE)GetProcAddress(GetModuleHandle(L"kernel32.dll"), "LoadLibraryA");
    InjectorCreateThread(proc, loadLibraryAddress, page);
    return 0;
}
